#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and 
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain 
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________
#
# Farmer: Generates scenario tree structure using callback
#
# Imports
#

from pyomo.core import *
import random

#
# Model
#

model = ConcreteModel()

#
# Parameters
#

model.CROPS = Set(initialize=['WHEAT','CORN','SUGAR_BEETS'])

model.TOTAL_ACREAGE = 500.0

model.PriceQuota = {'WHEAT':100000.0,'CORN':100000.0,'SUGAR_BEETS':6000.0}

model.SubQuotaSellingPrice = {'WHEAT':170.0,'CORN':150.0,'SUGAR_BEETS':36.0}

model.SuperQuotaSellingPrice = {'WHEAT':0.0,'CORN':0.0,'SUGAR_BEETS':10.0}

model.CattleFeedRequirement = {'WHEAT':200.0,'CORN':240.0,'SUGAR_BEETS':0.0}

model.PurchasePrice = {'WHEAT':238.0,'CORN':210.0,'SUGAR_BEETS':100000.0}

model.PlantingCostPerAcre = {'WHEAT':150.0,'CORN':230.0,'SUGAR_BEETS':260.0}

model.Yield = Param(model.CROPS, within=NonNegativeReals, initialize=0.0, mutable=True)

#
# Variables
#

model.DevotedAcreage = Var(model.CROPS, bounds=(0.0, model.TOTAL_ACREAGE))

model.QuantitySubQuotaSold = Var(model.CROPS, bounds=(0.0, None))
model.QuantitySuperQuotaSold = Var(model.CROPS, bounds=(0.0, None))
model.QuantityPurchased = Var(model.CROPS, bounds=(0.0, None))

#
# Constraints
#

def ConstrainTotalAcreage_rule(model):
    return sum_product(model.DevotedAcreage) <= model.TOTAL_ACREAGE

model.ConstrainTotalAcreage = Constraint(rule=ConstrainTotalAcreage_rule)

def EnforceCattleFeedRequirement_rule(model, i):
    return model.CattleFeedRequirement[i] <= (model.Yield[i] * model.DevotedAcreage[i]) + model.QuantityPurchased[i] - model.QuantitySubQuotaSold[i] - model.QuantitySuperQuotaSold[i]

model.EnforceCattleFeedRequirement = Constraint(model.CROPS, rule=EnforceCattleFeedRequirement_rule)

def LimitAmountSold_rule(model, i):
    return model.QuantitySubQuotaSold[i] + model.QuantitySuperQuotaSold[i] - (model.Yield[i] * model.DevotedAcreage[i]) <= 0.0

model.LimitAmountSold = Constraint(model.CROPS, rule=LimitAmountSold_rule)

def EnforceQuotas_rule(model, i):
    return (0.0, model.QuantitySubQuotaSold[i], model.PriceQuota[i])

model.EnforceQuotas = Constraint(model.CROPS, rule=EnforceQuotas_rule)

#
# Stage-specific cost computations
#

def ComputeFirstStageCost_rule(model):
    return sum_product(model.PlantingCostPerAcre, model.DevotedAcreage)
model.FirstStageCost = Expression(rule=ComputeFirstStageCost_rule)

def ComputeSecondStageCost_rule(model):
    expr = sum_product(model.PurchasePrice, model.QuantityPurchased)
    expr -= sum_product(model.SubQuotaSellingPrice, model.QuantitySubQuotaSold)
    expr -= sum_product(model.SuperQuotaSellingPrice, model.QuantitySuperQuotaSold)
    return expr
model.SecondStageCost = Expression(rule=ComputeSecondStageCost_rule)

#
# PySP Auto-generated Objective
#
# minimize: sum of StageCosts
#
# An active scenario objective equivalent to that generated by PySP is
# included here for informational purposes.
def total_cost_rule(model):
    return model.FirstStageCost + model.SecondStageCost
model.Total_Cost_Objective = Objective(rule=total_cost_rule, sense=minimize)

#
# Generate stochastic data as needed
#

# Number of scenarios
scenarios = 100

# Parameters for Normal distribution (mu, sigma)
RandomYield = {}
RandomYield['WHEAT'] = (2.5, 0.5)
RandomYield['CORN'] = (3.0, 0.4)
RandomYield['SUGAR_BEETS'] = (20.0, 4.0)

# Use a deterministic hash across python versions so we can semi-test
# this model
import hashlib
seed_random = lambda name: \
    random.seed(
        int(hashlib.sha512(name.encode()).hexdigest(), 16))

def pysp_scenario_tree_model_callback():
    from pyomo.pysp.scenariotree.tree_structure_model \
        import CreateConcreteTwoStageScenarioTreeModel

    st_model = CreateConcreteTwoStageScenarioTreeModel(scenarios)

    first_stage = st_model.Stages.first()
    second_stage = st_model.Stages.last()

    # First Stage
    st_model.StageCost[first_stage] = 'FirstStageCost'
    st_model.StageVariables[first_stage].add('DevotedAcreage[*]')

    # Second Stage
    st_model.StageCost[second_stage] = 'SecondStageCost'
    st_model.StageVariables[second_stage].add('QuantitySubQuotaSold[*]')
    st_model.StageVariables[second_stage].add('QuantitySuperQuotaSold[*]')
    st_model.StageVariables[second_stage].add('QuantityPurchased[*]')

    return st_model

def pysp_instance_creation_callback(scenario_name, node_names):
    
    instance = model.clone()

    seed_random(scenario_name)
    # Note: Must sort the iteration order so that
    #       data is generated deterministically
    #       across runs. Iteration over dict and set
    #       is very fickle in Python3.x
    for key in sorted(RandomYield.keys()):
        instance.Yield[key] = random.normalvariate(*RandomYield[key])

    return instance
